import psycopg2
import logging
from django.core.management.base import BaseCommand
from django.conf import settings


class Command(BaseCommand):
    help = __doc__

    def close_other_sessions(self, cursor, database_name):
        close_sessions_query = """
            SELECT pg_terminate_backend(pg_stat_activity.pid)
            FROM pg_stat_activity
            WHERE pg_stat_activity.datname = '%s';
        """ % database_name
        cursor.execute(close_sessions_query)

    def handle(self, *args, **options):
        if settings.ENVIRONMENT in ('staging', 'production'):
            raise Exception('Stop it!')

        # Prepare connection and db cursor
        dbinfo = settings.DATABASES.get('default')
        conn_params = {
            'database': 'template1',
            'user': dbinfo.get('USER'),
            'password': dbinfo.get('PASSWORD'),
            'host': dbinfo.get('HOST'),
            'port': dbinfo.get('PORT'),
        }
        connection = psycopg2.connect(**conn_params)
        connection.set_isolation_level(0)  # autocommit false
        cursor = connection.cursor()

        # Forcefully close other still connected sessions
        database_name = dbinfo.get('NAME')
        self.close_other_sessions(cursor, database_name)

        # Drop existing database
        drop_query = "DROP DATABASE \"%s\";" % database_name
        logging.info('Executing... "%s"', drop_query)
        try:
            cursor.execute(drop_query)
        except psycopg2.ProgrammingError as e:
            logging.exception("Error: %s", str(e))

        # Create the dropped database (clean)
        create_query = "CREATE DATABASE \"%s\"" % database_name
        create_query += " WITH OWNER = \"%s\" " % dbinfo.get('USER')
        create_query += " ENCODING = 'UTF8';"

        logging.info('Executing... "%s"', create_query)
        cursor.execute(create_query)
